% Counter-factual with lock-down

clc
clear all
close all

% calibrated beta and d_I
beta=0.15246;
d_I=0.01634;

w=ones(26,2); % no wage differentials
load moments_from_Seoul_survey.mat
load distance_GIS.mat
load calibrated_E.mat
load initial_infection.mat

epsilon_weekday=4.1642;
epsilon_weekend=4.9144;
kappa=0.0339;

H_R(:,1)=H_R_young; % young
H_R(:,2)=H_R_old;   % old

ncity=25;

gamma=1/18; % 18 days to be recovered
tau=[1/8.5 1/10.2]; % days to be quaurantined
death_rate=[0.0021 0.0273]; % fatality rate

delta_j=zeros(4,2); % no disclosure

delta_t=zeros(2,2);
lockdown=0.22; % lockdown coefficient (0.3 means that we put 30% to home sector and allow 70% to choose commuting locations)

start_day=280;
end_day=380; % 100 days

global w d_GIS ncity epsilon_weekday epsilon_weekend kappa E_weekday E_weekend H_R initial_infection gamma tau death_rate beta d_I delta_j delta_t

time=1000; period=[1:1:time];

%% Actual Simulation
% setup dates: starting from Jan 30
Day = datetime(2020,1,30) + caldays(0:time);

% Initial Values
    Qtime(1,:,1)=initial_infection'; % Initial detected, as of January 30 (all looks exogenous)
    % city 8(1), 13(1), 15(1), 16(1) // 4 the detected
    Qtime(1,:,2)=zeros(1,25);
    Itime(1,:,1)=9*initial_infection';
    Itime(1,:,2)=zeros(1,25);
    for a=1:2
        for d1=1:ncity
            Ntime(1,d1,a)=H_R(d1,a);
            Stime(1,d1,a)=Ntime(1,d1,a)-Itime(1,d1,a)-Qtime(1,d1,a);
            Rtime(1,d1,a)=Ntime(1,d1,a)-Itime(1,d1,a)-Stime(1,d1,a)-Qtime(1,d1,a); % zero
            Ftime(1,d1,a)=death_rate(a)*Rtime(1,d1,a);
            Dtime(1,d1,a)=0;
            cum_Dtime(1,d1)=Qtime(1,d1,1)+Qtime(1,d1,2);
            cum_Dtime_14(1,d1)=Qtime(1,d1,1)+Qtime(1,d1,2);
        end
    end

% run simulation
for t=2:time
    [DayNumber,DayName] = weekday(Day(t));
    if DayNumber>=2 && DayNumber<=6 % weekdays
            [pi(t,:,:,:) d_ija(t,:,:,:)]=Calculate_pi_function(cum_Dtime_14(t-1,:));
            if t<=start_day || t>end_day % before lockdown or after lockdown
            else % lockdown period
            for a=1:2
            for d1=1:25
                for d2=1:25
                    pi(t,d1,d2,a)=(1-lockdown)*pi(t,d1,d2,a); % remove lockdown fraction of people from everywhere
                end
                pi(t,d1,26,a)=(1-lockdown)*pi(t,d1,26,a)+lockdown; % put them in home sector
            end
            end
            end
                for d1=1:ncity
                    for a=1:2
                    S_dot(t,d1,a)=0;
                    for d2=1:ncity
                        S_dot(t,d1,a)=S_dot(t,d1,a)-beta*(sum(pi(t,:,d2,1).*Itime(t-1,:,1))+sum(pi(t,:,d2,2).*Itime(t-1,:,2)))/(sum(pi(t,:,d2,1).*(Ntime(t-1,:,1)-Qtime(t-1,:,1)))+sum(pi(t,:,d2,2).*(Ntime(t-1,:,2)-Qtime(t-1,:,2))))*(pi(t,d1,d2,a)*Stime(t-1,d1,a));
                    end
                    Stime(t,d1,a) =  Stime(t-1,d1,a) + S_dot(t,d1,a);
                    Itime(t,d1,a) =  Itime(t-1,d1,a)*(1 - gamma) - S_dot(t,d1,a) - d_I*Itime(t-1,d1,a);
                    Qtime(t,d1,a) =  Qtime(t-1,d1,a)*(1 - tau(a)) + d_I*Itime(t-1,d1,a);
                    Rtime(t,d1,a) =  Rtime(t-1,d1,a) + Itime(t-1,d1,a)*gamma + tau(a)*Qtime(t-1,d1,a);
                    Ftime(t,d1,a) =  death_rate(a)*Rtime(t,d1,a);
                    Ntime(t,d1,a) =  H_R(d1,a) - death_rate(a)*Rtime(t,d1,a);
                    Dtime(t,d1,a) = d_I*Itime(t-1,d1,a);
                    end
                    cum_Dtime(t,d1) = cum_Dtime(t-1,d1) + d_I*Itime(t-1,d1,1) + d_I*Itime(t-1,d1,2); % adding newly detected cases
                    cum_Dtime_14(t,d1) = sum(Dtime(max(1,t-13):t,d1,1))+sum(Dtime(max(1,t-13):t,d1,2));
                end
    else % weekends
            [pi(t,:,:,:) d_ija(t,:,:,:)]=Calculate_pi_weekend_function(cum_Dtime_14(t-1,:));
            if t<=start_day || t>end_day % before lockdown or after lockdown
            else % lockdown period
            for a=1:2
            for d1=1:25
                for d2=1:25
                    pi(t,d1,d2,a)=(1-lockdown)*pi(t,d1,d2,a); % remove lockdown fraction of people from everywhere
                end
                pi(t,d1,26,a)=(1-lockdown)*pi(t,d1,26,a)+lockdown; % put them in home sector
            end
            end
            end
                for d1=1:ncity
                    for a=1:2
                    S_dot(t,d1,a)=0;
                    for d2=1:ncity
                        S_dot(t,d1,a)=S_dot(t,d1,a)-beta*(sum(pi(t,:,d2,1).*Itime(t-1,:,1))+sum(pi(t,:,d2,2).*Itime(t-1,:,2)))/(sum(pi(t,:,d2,1).*(Ntime(t-1,:,1)-Qtime(t-1,:,1)))+sum(pi(t,:,d2,2).*(Ntime(t-1,:,2)-Qtime(t-1,:,2))))*(pi(t,d1,d2,a)*Stime(t-1,d1,a));
                    end
                    Stime(t,d1,a) =  Stime(t-1,d1,a) + S_dot(t,d1,a);
                    Itime(t,d1,a) =  Itime(t-1,d1,a)*(1 - gamma) - S_dot(t,d1,a) - d_I*Itime(t-1,d1,a);
                    Qtime(t,d1,a) =  Qtime(t-1,d1,a)*(1 - tau(a)) + d_I*Itime(t-1,d1,a);
                    Rtime(t,d1,a) =  Rtime(t-1,d1,a) + Itime(t-1,d1,a)*gamma + tau(a)*Qtime(t-1,d1,a);
                    Ftime(t,d1,a) =  death_rate(a)*Rtime(t,d1,a);
                    Ntime(t,d1,a) =  H_R(d1,a) - death_rate(a)*Rtime(t,d1,a);
                    Dtime(t,d1,a) = d_I*Itime(t-1,d1,a);
                    end
                    cum_Dtime(t,d1) = cum_Dtime(t-1,d1) + d_I*Itime(t-1,d1,1) + d_I*Itime(t-1,d1,2); % adding newly detected cases
                    cum_Dtime_14(t,d1) = sum(Dtime(max(1,t-13):t,d1,1))+sum(Dtime(max(1,t-13):t,d1,2));
                end
    end
end

day=period;
Pop=sum(H_R);

common_z_weekday=draw_frechet(E_weekday,epsilon_weekday,10^4);
common_z_weekend=draw_frechet(E_weekend,epsilon_weekend,10^4);

global common_z_weekday common_z_weekend

for t=1:time
    [DayNumber,DayName] = weekday(Day(t));
    if DayNumber>=2 && DayNumber<=6 % weekdays
        if t<=start_day || t>end_day
            utility_weekdays(t,:,:)=Calculate_utility_per_capita_function(cum_Dtime_14(t,:),squeeze(Qtime(t,:,:)),squeeze(Rtime(t,:,:)));
        else
            utility_weekdays(t,:,:)=Calculate_utility_per_capita_lockdown_function(cum_Dtime_14(t,:),squeeze(Qtime(t,:,:)),squeeze(Rtime(t,:,:)),lockdown);
        end
    else
        if t<=start_day || t>end_day
            utility_weekends(t,:,:)=Calculate_utility_per_capita_weekend_function(cum_Dtime_14(t,:),squeeze(Qtime(t,:,:)),squeeze(Rtime(t,:,:)));
        else
            utility_weekends(t,:,:)=Calculate_utility_per_capita_weekend_lockdown_function(cum_Dtime_14(t,:),squeeze(Qtime(t,:,:)),squeeze(Rtime(t,:,:)),lockdown);
        end
    end
    t
end

utility_weekdays_sum(:,1)=sum(utility_weekdays(:,:,1),2);
utility_weekdays_sum(:,2)=sum(utility_weekdays(:,:,2),2);
size_weekdays=size(utility_weekdays_sum);
for a=1:2
for t=1:size_weekdays(1)
    if utility_weekdays_sum(t,a)==0
        utility_weekdays_sum(t,a)=NaN;
    end
end
end
utility_weekdays_sum_norm=log(utility_weekdays_sum./utility_weekdays_sum(1,:));

utility_weekends_sum(:,1)=sum(utility_weekends(:,:,1),2);
utility_weekends_sum(:,2)=sum(utility_weekends(:,:,2),2);
size_weekends=size(utility_weekends_sum);
for a=1:2
for t=1:size_weekends(1)
    if utility_weekends_sum(t,a)==0
        utility_weekends_sum(t,a)=NaN;
    end
end
end
utility_weekends_sum_norm=log(utility_weekends_sum./utility_weekends_sum(3,:)); % first weekend, day 3

utility_weekdays_sum_norm_avg=utility_weekdays_sum_norm(:,1)*(Pop(1)/(Pop(1)+Pop(2)))+utility_weekdays_sum_norm(:,2)*(Pop(2)/(Pop(1)+Pop(2)));
utility_weekends_sum_norm_avg=utility_weekends_sum_norm(:,1)*(Pop(1)/(Pop(1)+Pop(2)))+utility_weekends_sum_norm(:,2)*(Pop(2)/(Pop(1)+Pop(2)));

for t=2:time
    [DayNumber,DayName] = weekday(Day(t));
    if DayNumber>=2 && DayNumber<=6 % weekdays
        utility_sum_norm_avg(t,1)=utility_weekdays_sum_norm_avg(t);
        utility_sum_norm_young(t,1)=utility_weekdays_sum_norm(t,1);
        utility_sum_norm_old(t,1)=utility_weekdays_sum_norm(t,2);
    else
        utility_sum_norm_avg(t,1)=utility_weekends_sum_norm_avg(t);
        utility_sum_norm_young(t,1)=utility_weekends_sum_norm(t,1);
        utility_sum_norm_old(t,1)=utility_weekends_sum_norm(t,2);
    end
end

%% Save for cnt
total_cum_Dtime=sum(cum_Dtime,2);
save ('Output/cnt_lockdown.mat', 'day', 'Stime', 'Itime', 'Qtime', 'Rtime', 'Ftime', 'total_cum_Dtime', 'utility_sum_norm_avg', 'utility_sum_norm_young','utility_sum_norm_old')

